#include "png_data.h"

namespace ux_png {


void PngData::InitInfoObject(int channels, int color_type, int bit_depth) {
  channels_ = channels;
  color_type_ = color_type;
  bit_depth_ = bit_depth;
  have_alpha_ = (color_type & PNG_COLOR_MASK_ALPHA) != 0;
}

PngData::~PngData() {

}

int PngData::Channels() const {
  return channels_;
}

int PngData::ColorType() const {
  return color_type_;
}

int PngData::Width() const {
  if (!mp_.empty()) return mp_[0].size();
  else return 0;
}

int PngData::Height() const {
  return mp_.size();
}

int PngData::BitDepth() const {
  return bit_depth_;
}

int PngData::HaveAlpha() const {
  return have_alpha_;
}

int PngData::LeftPut(const PngData &other) {
  if (mp_.size() != other.mp_.size()) {
    printf("height is not equal, this.height=%d\n, other height=%d\n", Height(), other.Height());
    return -1;
  }
  for (int i = 0; i < Height(); i++) {
    for (int j = 0; j < other.Width(); j++) {
      mp_[i].push_back(other.mp_[i][j]);
    }
  }
  return 0;
}

void PngData::SetColorType(int new_color_type) {
  color_type_ = new_color_type;
  channels_ = ColorType2Channels(color_type_);
}

int PngData::Rgb2Gray() {
  if (ColorType() == PNG_COLOR_TYPE_GRAY ||
      ColorType() == PNG_COLOR_TYPE_GA) {
    printf("this png is already gray\n");
    return -1;
  }
  int new_color_type;
  if (ColorType() == PNG_COLOR_TYPE_RGB) {
    new_color_type = PNG_COLOR_TYPE_GRAY;
  }
  else if (ColorType() == PNG_COLOR_TYPE_RGBA) {
    new_color_type = PNG_COLOR_TYPE_GA;
  } else {
    printf("error color_type %d\n", color_type_);
    return -1;
  }
  for (auto & pixel_v : mp_) {
    for (auto & p : pixel_v) {
      p = p.ToGray();
    }
  }
  SetColorType(new_color_type);
  return 0;
}

int ReadPngDataFromImgFile(const char * img_name, PngData* png_data)  {
  png_structp png_ptr;
  png_infop info_ptr;
  FILE *fp;
  if (!png_data) {
    printf("png_data is null\n");
    return -1;
  }
  if (!IsPNG(img_name)) {
    printf("file is not png ...\n");
    return -1;
  }
  fp = fopen(img_name, "rb");

  png_ptr  = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
  info_ptr = png_create_info_struct(png_ptr);

  png_init_io(png_ptr, fp);
  png_read_png(png_ptr, info_ptr, PNG_TRANSFORM_EXPAND, nullptr);

  int channels, color_type, bit_depth, width, height;
  channels = png_get_channels(png_ptr, info_ptr);
  color_type = png_get_color_type(png_ptr, info_ptr);
  bit_depth = png_get_bit_depth(png_ptr, info_ptr);
  width = png_get_image_width(png_ptr, info_ptr);
  height  = png_get_image_height(png_ptr, info_ptr);

  printf("channels = %d color_type = %d bit_depth = %d width = %d height = %d\n",
         channels, color_type, bit_depth, width, height);

  png_data->InitInfoObject(channels, color_type, bit_depth);
  auto row_pointers = png_get_rows(png_ptr, info_ptr);
  png_data->mp_.resize(height, std::vector<Pixel>(width));
  for (int i = 0; i < height; i++) {
    for (int j = 0; j < width; j++) {
      auto block = channels;
      for (int pixel_ind = 0, k = block * (j + 1) - 1; k >= block * j; k--, pixel_ind++) {
        png_data->mp_[i][j].set(row_pointers[i][k], pixel_ind);
      }
      png_data->mp_[i][j].color_type_ = png_data->ColorType();
    }
  }
  png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
  fclose(fp);
  return 0;
}

int ReadAndCheckPNGFromImgFile(const char * img_name, PngData* png_data) {
  auto is_png = IsPNG(img_name);
  if (!is_png) {
    printf("%s is not a png\n", img_name);
    return -1;
  }
  return ReadPngDataFromImgFile(img_name, png_data);
}

int WritePngDataToImgFile(const char * img_name, const PngData* png_data) {
  png_structp png_ptr;
  png_infop info_ptr;
  FILE* fp;
  if (!png_data) {
    printf("png_data is null\n");
    return -1;
  }
  fp = fopen(img_name, "wb");
  if (!fp) {
    printf("open file failed\n");
    return -1;
  }
  png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
  if (!png_ptr) {
    printf("png_ptr create err\n");
    return -1;
  }
  info_ptr = png_create_info_struct(png_ptr);
  if (!info_ptr) {
    printf("info_ptr create err\n");
    return -1;
  }
  png_init_io(png_ptr, fp);
  png_set_IHDR(png_ptr, info_ptr,
          png_data->Width(), png_data->Height(), png_data->BitDepth(), png_data->ColorType(),
          PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
  png_write_info(png_ptr, info_ptr);
  for (int i = 0; i < png_data->Height(); i++) {
    int block = png_data->Channels();
    auto row_ptr = new png_byte[block * png_data->Width()];
    for (int j = 0; j < png_data->Width(); j++) {
      for (int k = j * block, pixel_id = block - 1; pixel_id >= 0; pixel_id--, k++) {
        row_ptr[k] = png_data->mp_[i][j].get(pixel_id);
      }
    }
    png_write_row(png_ptr, row_ptr);
    delete[] row_ptr;
  }
  png_write_end(png_ptr, nullptr);
  png_destroy_write_struct(&png_ptr, &info_ptr);
  fclose(fp);
  return 0;
}

Pixel Pixel::ToGray() {
  auto len = ColorType2Channels(color_type_);
  if (len < 3) return {};
  auto gray = static_cast<unsigned char>(0.299*_[0] + 0.587*_[1] + 0.114*_[2]);
  Pixel p;
  p.set(gray, 0);
  p.color_type_ = PNG_COLOR_TYPE_GRAY;
  if (len == 4) {
    p.set(_[3], 1);
    p.color_type_ = PNG_COLOR_TYPE_GRAY_ALPHA;
  }
  return p;
}

}